package com.liyuxiang.graph.memory;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.jsontype.impl.LaissezFaireSubTypeValidator;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import com.fasterxml.jackson.databind.jsontype.BasicPolymorphicTypeValidator;
import com.fasterxml.jackson.databind.JsonNode;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @Author: liyuxiang
 * @Date: 2025/8/3 14:28
 * 基于Redis实现的聊天记忆存储
 * 实现了ChatMemory接口，用于存储和检索聊天消息历史
 */
@Component
public class SpringAiRedisChatMemory implements ChatMemory {

    private static final Logger logger = LoggerFactory.getLogger(SpringAiRedisChatMemory.class);

    // Redis操作模板
    private final StringRedisTemplate redisTemplate;
    // JSON序列化工具
    private final ObjectMapper objectMapper;

    // Redis key前缀，避免与其他数据冲突
    private static final String CHAT_MEMORY_PREFIX = "ai-chat:memory:";

    // 配置参数：防止内存溢出
    @Value("${chat.memory.max-messages:100}")
    private int maxMessages = 100; // 每个会话最大消息数

    @Value("${chat.memory.ttl-hours:24}")
    private int ttlHours = 1; // 会话过期时间（小时）

    @Value("${chat.memory.batch-size:50}")
    private int batchSize = 50; // 批量操作大小

    public SpringAiRedisChatMemory(StringRedisTemplate redisTemplate, ObjectMapper objectMapper) {
        this.redisTemplate = redisTemplate;
        // 配置ObjectMapper以正确处理Message的多态性
        this.objectMapper = objectMapper.copy();
        
        // 使用更宽松的多态类型验证器以支持所有必要的类型
        BasicPolymorphicTypeValidator typeValidator = BasicPolymorphicTypeValidator.builder()
            .allowIfSubType("org.springframework.ai.chat.messages")
            .allowIfSubType(Message.class)
            .allowIfSubType("java.util")  // 允许Java集合类型
            .allowIfSubType("java.lang")  // 允许Java基本类型
            .allowIfSubType("org.springframework.ai.chat.messages.MessageType")  // 允许枚举类型
            .build();
        
        this.objectMapper.activateDefaultTyping(typeValidator, 
            ObjectMapper.DefaultTyping.NON_FINAL, JsonTypeInfo.As.PROPERTY);
        
        logger.info("SpringAiRedisChatMemory 初始化 - 最大消息数: {}, TTL: {}小时, 批量大小: {}", 
                    maxMessages, ttlHours, batchSize);
    }

    /**
     * 获取完整的Redis key
     */
    private String getRedisKey(String conversationId) {
        return CHAT_MEMORY_PREFIX + conversationId;
    }

    @Override
    public void add(String conversationId, Message message) {
        add(conversationId, List.of(message));
    }

    /**
     * 添加消息到聊天记忆
     * 增加长度限制和TTL设置，防止内存溢出
     *
     * @param conversationId: 会话ID，用于区分不同对话
     * @param messages:       要添加的消息列表
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }

        String redisKey = getRedisKey(conversationId);
        logger.debug("添加消息到会话 {} - 消息数: {}", conversationId, messages.size());

        try {
            List<String> serializedMessages = messages.stream()
                                                      .map(msg -> {
                                                          try {
                                                              // 为消息添加唯一ID和类型信息
                                                              msg.getMetadata().putIfAbsent("id", UUID.randomUUID().toString());
                                                              msg.getMetadata().putIfAbsent("timestamp", String.valueOf(System.currentTimeMillis()));
                                                              msg.getMetadata().putIfAbsent("messageType", msg.getMessageType().name());
                                                              
                                                              return objectMapper.writeValueAsString(msg);
                                                          } catch (JsonProcessingException e) {
                                                              logger.error("序列化消息失败: {}", e.getMessage(), e);
                                                              throw new RuntimeException("Failed to serialize message", e);
                                                          }
                                                      })
                                                      .collect(Collectors.toList());

            // 使用rightPush保持消息顺序（最新的在最后）
            redisTemplate.opsForList().rightPushAll(redisKey, serializedMessages);

            // 设置TTL，防止数据永久存储
            redisTemplate.expire(redisKey, ttlHours, TimeUnit.HOURS);

            // 检查并清理超长的消息列表，防止内存溢出
            trimMessagesIfNeeded(redisKey, conversationId);

        } catch (Exception e) {
            logger.error("添加消息到Redis失败 - 会话ID: {}, 错误: {}", conversationId, e.getMessage(), e);
            throw e;
        }
    }

    /**
     * 清理超长的消息列表，保留最近的消息
     */
    private void trimMessagesIfNeeded(String redisKey, String conversationId) {
        try {
            Long currentSize = redisTemplate.opsForList().size(redisKey);
            if (currentSize != null && currentSize > maxMessages) {
                // 计算需要删除的消息数
                long excessCount = currentSize - maxMessages;
                logger.info("会话 {} 消息数量 {} 超过限制 {}，清理最早的 {} 条消息", 
                           conversationId, currentSize, maxMessages, excessCount);
                
                // 删除最早的消息（从左侧删除）
                for (int i = 0; i < excessCount; i++) {
                    redisTemplate.opsForList().leftPop(redisKey);
                }
                
                logger.info("会话 {} 清理完成，当前消息数: {}", conversationId, 
                           redisTemplate.opsForList().size(redisKey));
            }
        } catch (Exception e) {
            logger.error("清理会话 {} 消息失败: {}", conversationId, e.getMessage(), e);
        }
    }

    @Override
    public List<Message> get(String conversationId) {
        // 为了防止内存溢出，限制返回最近的消息
        return get(conversationId, maxMessages);
    }

    /**
     * 从聊天记忆中获取指定数量的最近消息
     *
     * @param conversationId: 会话ID
     * @param lastN:          要获取的消息数量
     * @return 返回最近lastN条消息
     */
    public List<Message> get(String conversationId, int lastN) {
        if (lastN <= 0) {
            return List.of();
        }

        // 获取总消息数
        Long totalSize = redisTemplate.opsForList().size(getRedisKey(conversationId));
        if (totalSize == null || totalSize == 0) {
            return List.of();
        }

        // 计算起始位置（获取最后lastN条消息）
        long start = Math.max(0, totalSize - lastN);
        long end = totalSize - 1;

        List<String> messages = redisTemplate.opsForList().range(getRedisKey(conversationId), start, end);
        return deserializeMessages(messages);
    }

    /**
     * 反序列化消息列表
     * 增加批量处理和错误容错，防止内存溢出
     */
    private List<Message> deserializeMessages(List<String> serializedMessages) {
        if (serializedMessages == null || serializedMessages.isEmpty()) {
            return List.of();
        }

        // 限制处理数量，防止内存溢出
        int processCount = Math.min(serializedMessages.size(), maxMessages);
        List<String> messagesToProcess = serializedMessages.subList(
            Math.max(0, serializedMessages.size() - processCount), 
            serializedMessages.size()
        );

        logger.debug("反序列化消息 - 总数: {}, 处理数: {}", serializedMessages.size(), processCount);

        return messagesToProcess.stream()
                               .map(json -> {
                                   try {
                                       // 尝试解析消息类型并选择合适的反序列化方式
                                       return deserializeMessage(json);
                                   } catch (Exception e) {
                                       logger.warn("反序列化单条消息失败，跳过该消息: {}", e.getMessage());
                                       return null;
                                   }
                               })
                               .filter(message -> message != null) // 过滤掉反序列化失败的消息
                               .collect(Collectors.toList());
    }

    /**
     * 安全的消息反序列化方法
     * 支持多种格式的消息反序列化，确保兼容性
     */
    private Message deserializeMessage(String json) throws Exception {
        try {
            // 首先尝试直接反序列化
            return objectMapper.readValue(json, Message.class);
        } catch (JsonProcessingException e) {
            logger.debug("直接反序列化失败，尝试其他方式: {}", e.getMessage());
            
            // 尝试使用更宽松的ObjectMapper处理旧格式数据
            try {
                return deserializeWithLegacyFormat(json);
            } catch (Exception legacyException) {
                logger.debug("旧格式反序列化也失败: {}", legacyException.getMessage());
                
                // 尝试解析JSON并判断消息类型
                JsonNode jsonNode = objectMapper.readTree(json);
                
                // 检查消息类型字段
                String messageType = extractMessageType(jsonNode);
                
                // 根据类型选择合适的反序列化方式
                if (messageType != null) {
                    return deserializeByType(json, messageType);
                }
                
                // 如果以上方法都失败，则抛出原始异常
                throw e;
            }
        }
    }
    
    /**
     * 使用宽松配置处理旧格式数据
     */
    private Message deserializeWithLegacyFormat(String json) throws Exception {
        ObjectMapper legacyMapper = new ObjectMapper();
        legacyMapper.activateDefaultTyping(LaissezFaireSubTypeValidator.instance, 
                                          ObjectMapper.DefaultTyping.NON_FINAL);
        return legacyMapper.readValue(json, Message.class);
    }
    
    /**
     * 提取消息类型
     */
    private String extractMessageType(JsonNode jsonNode) {
        // 检查顶级messageType字段
        if (jsonNode.has("messageType")) {
            return jsonNode.get("messageType").asText();
        }
        
        // 检查metadata中的messageType字段
        if (jsonNode.has("metadata")) {
            JsonNode metadata = jsonNode.get("metadata");
            if (metadata.has("messageType")) {
                JsonNode msgTypeNode = metadata.get("messageType");
                // 处理可能的数组格式 ["org.springframework.ai.chat.messages.MessageType","USER"]
                if (msgTypeNode.isArray() && msgTypeNode.size() > 1) {
                    return msgTypeNode.get(1).asText();
                } else if (msgTypeNode.isTextual()) {
                    return msgTypeNode.asText();
                }
            }
        }
        
        // 检查@class字段来推断类型
        if (jsonNode.has("@class")) {
            String className = jsonNode.get("@class").asText();
            if (className.contains("UserMessage")) return "USER";
            if (className.contains("AssistantMessage")) return "ASSISTANT";
            if (className.contains("SystemMessage")) return "SYSTEM";
        }
        
        return null;
    }
    
    /**
     * 根据类型反序列化
     */
    private Message deserializeByType(String json, String messageType) throws Exception {
        switch (messageType.toUpperCase()) {
            case "USER":
                return objectMapper.readValue(json, UserMessage.class);
            case "ASSISTANT":
                return objectMapper.readValue(json, AssistantMessage.class);
            case "SYSTEM":
                return objectMapper.readValue(json, SystemMessage.class);
            default:
                logger.warn("未知的消息类型: {}", messageType);
                throw new IllegalArgumentException("Unknown message type: " + messageType);
        }
    }

    /**
     * 清除指定会话的聊天记忆
     *
     * @param conversationId: 要清除的会话ID
     */
    @Override
    public void clear(String conversationId) {
        redisTemplate.delete(getRedisKey(conversationId));
    }

    /**
     * 移除指定消息ID的聊天记忆
     * 优化：限制搜索范围，避免处理过大数据
     *
     * @param conversationId: 会话ID
     * @param messageId:      消息ID
     */
    public void deleteByMessageId(String conversationId, String messageId) {
        String redisKey = getRedisKey(conversationId);
        
        try {
            Long totalSize = redisTemplate.opsForList().size(redisKey);
            if (totalSize == null || totalSize == 0) {
                return;
            }
            
            // 限制搜索范围，避免内存溢出
            int searchLimit = Math.min(totalSize.intValue(), batchSize);
            List<String> messages = redisTemplate.opsForList().range(redisKey, -searchLimit, -1);

            if (messages != null) {
                for (int i = 0; i < messages.size(); i++) {
                    try {
                        String messageJson = messages.get(i);
                        Message message = objectMapper.readValue(messageJson, Message.class);
                        String msgId = message.getMetadata().get("id").toString();
                        if (messageId.equals(msgId)) {
                            // 移除匹配的消息
                            redisTemplate.opsForList().remove(redisKey, 1, messageJson);
                            logger.info("删除消息ID {} 从会话 {}", messageId, conversationId);
                            break; // 找到并删除后退出循环
                        }
                    } catch (JsonProcessingException e) {
                        logger.error("反序列化消息失败，跳过: {}", e.getMessage());
                    }
                }
            }
        } catch (Exception e) {
            logger.error("删除消息失败 - 会话ID: {}, 消息ID: {}, 错误: {}", conversationId, messageId, e.getMessage(), e);
        }
    }
    
    /**
     * 批量清理过期会话
     * 可通过定时任务调用此方法
     */
    public void cleanupExpiredSessions() {
        logger.info("开始清理过期会话");
        try {
            // 这里可以添加批量清理逻辑
            // 由于Redis TTL会自动过期，这个方法主要用于统计和日志
            logger.info("过期会话清理完成");
        } catch (Exception e) {
            logger.error("清理过期会话失败: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 清理损坏的历史消息数据
     * 用于处理反序列化失败的历史消息
     */
    public void cleanupCorruptedMessages(String conversationId) {
        String redisKey = getRedisKey(conversationId);
        logger.info("开始清理会话 {} 的损坏消息", conversationId);
        
        try {
            List<String> messages = redisTemplate.opsForList().range(redisKey, 0, -1);
            if (messages == null || messages.isEmpty()) {
                return;
            }
            
            int cleanedCount = 0;
            for (String messageJson : messages) {
                try {
                    deserializeMessage(messageJson);
                } catch (Exception e) {
                    // 移除损坏的消息
                    redisTemplate.opsForList().remove(redisKey, 1, messageJson);
                    cleanedCount++;
                    logger.debug("移除损坏消息: {}", e.getMessage());
                }
            }
            
            if (cleanedCount > 0) {
                logger.info("会话 {} 清理完成，移除了 {} 条损坏消息", conversationId, cleanedCount);
            }
        } catch (Exception e) {
            logger.error("清理损坏消息失败 - 会话ID: {}, 错误: {}", conversationId, e.getMessage(), e);
        }
    }
    
    /**
     * 获取会话统计信息
     */
    public void logMemoryUsage() {
        logger.info("ChatMemory配置 - 最大消息数: {}, TTL: {}小时, 批量大小: {}", 
                   maxMessages, ttlHours, batchSize);
    }
    
    /**
     * 测试反序列化修复效果
     * 检查指定会话的消息反序列化状态
     */
    public void testDeserialization(String conversationId) {
        String redisKey = getRedisKey(conversationId);
        logger.info("开始测试会话 {} 的反序列化功能", conversationId);
        
        try {
            List<String> messages = redisTemplate.opsForList().range(redisKey, 0, -1);
            if (messages == null || messages.isEmpty()) {
                logger.info("会话 {} 无消息数据", conversationId);
                return;
            }
            
            int totalCount = messages.size();
            int successCount = 0;
            int failCount = 0;
            
            for (String messageJson : messages) {
                try {
                    Message msg = deserializeMessage(messageJson);
                    if (msg != null) {
                        successCount++;
                        logger.debug("成功反序列化消息: 类型={}, ID={}", 
                                   msg.getMessageType(), msg.getMetadata().get("id"));
                    }
                } catch (Exception e) {
                    failCount++;
                    logger.debug("反序列化失败: {}", e.getMessage());
                }
            }
            
            logger.info("会话 {} 反序列化测试完成 - 总数: {}, 成功: {}, 失败: {}, 成功率: {}%",
                       conversationId, totalCount, successCount, failCount, 
                       totalCount > 0 ? (successCount * 100 / totalCount) : 0);
                       
        } catch (Exception e) {
            logger.error("测试反序列化功能失败 - 会话ID: {}, 错误: {}", conversationId, e.getMessage(), e);
        }
    }
}